import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split

m = 10
x1 = np.arange(m)
x2 = x1 * 2
x3 = x1 * 3
x4 = x1 * 4
x = np.c_[x1, x2, x3, x4]
x = pd.DataFrame(x)
print(x)

offset_orig = np.arange(m)
rate = 0.7
offset_train, offset_test = train_test_split(offset_orig, train_size=rate, random_state=666)
print(offset_train)
print(offset_test)
if rate < 0.5:
    offset_set_train = set(offset_train)
    idx_train = [i in offset_set_train for i in offset_orig]
    idx_test = np.invert(idx_train)
else:
    offset_set_test = set(offset_test)
    idx_test = [i in offset_set_test for i in offset_orig]
    idx_train = np.invert(idx_test)
print(idx_train)
print(idx_test)
x_train = x[idx_train]  # later
x_test = x[idx_test]
print(x_train)
print(x_test)
